package com.piece.core.db.context;

import com.alibaba.druid.pool.DruidDataSource;
import com.alibaba.druid.stat.DruidDataSourceStatManager;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import javax.sql.DataSource;
import java.sql.DriverManager;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

/**
 * 动态数据源
 */
public class DynamicDataSource extends AbstractRoutingDataSource {

    private Map<Object, Object> dynamicTargetDataSources = new HashMap<>();

    public DynamicDataSource(DataSource defaultTargetDataSource, Map<Object, Object> targetDataSources) {
        super.setDefaultTargetDataSource(defaultTargetDataSource);
        this.setTargetDataSources(targetDataSources);
    }

    @Override
    protected Object determineCurrentLookupKey() {
        return DynamicDataSourceContextHolder.getDateSourceType();
    }

    @Override
    public void setTargetDataSources(Map<Object, Object> targetDataSources) {
        this.dynamicTargetDataSources.putAll(targetDataSources);
        super.setTargetDataSources(this.dynamicTargetDataSources);
        super.afterPropertiesSet();
    }

    public DruidDataSource getTargetDataSources(String key) {
        return (DruidDataSource) this.dynamicTargetDataSources.get(key);
    }

    /**
     * 创建数据源
     *
     * @param key
     * @param driveClass
     * @param url
     * @param username
     * @param password
     * @return
     */
    public boolean createDataSource(String key, String driveClass, String url, String username, String password) {
        try {
            try {
                // 排除连接不上的错误
                Class.forName(driveClass);
                DriverManager.getConnection(url, username, password);
            } catch (Exception e) {
                return false;
            }

            DruidDataSource druidDataSource = new DruidDataSource();
            druidDataSource.setName(key);
            druidDataSource.setDriverClassName(driveClass);
            druidDataSource.setUrl(url);
            druidDataSource.setUsername(username);
            druidDataSource.setPassword(password);
            druidDataSource.setMaxWait(60000);
            druidDataSource.setFilters("stat");
            DataSource createDataSource = druidDataSource;
            druidDataSource.init();
            Map<Object, Object> dynamicTargetDataSources = new HashMap<>();
            // 加入map
            dynamicTargetDataSources.put(key, createDataSource);
            // 将map赋值给父类的TargetDataSources
            setTargetDataSources(dynamicTargetDataSources);
            return true;
        } catch (Exception e) {
            return false;
        }
    }

    /**
     * 删除数据源
     *
     * @param key
     * @return
     */
    public boolean delDataSource(String key) {
        if (this.dynamicTargetDataSources.containsKey(key)) {
            Set<DruidDataSource> druidDataSourceInstances = DruidDataSourceStatManager.getDruidDataSourceInstances();
            for (DruidDataSource druidDataSource : druidDataSourceInstances) {
                if (key.equals(druidDataSource.getName())) {
                    DruidDataSourceStatManager.removeDataSource(druidDataSource);
                    // 将map赋值给父类的TargetDataSources
                    this.dynamicTargetDataSources.remove(druidDataSource);
                    setTargetDataSources(this.dynamicTargetDataSources);
                    return true;
                }
            }
            return false;
        } else {
            return false;
        }
    }
}
